from Agent import Agent, global_utility
import numpy as np
import random

global days, num_nights, num_agents, discount_rate


# 初始化
def setup():
    global days, num_nights, num_agents, discount_rate
    #     reward_function, actions

    # days = int(input("How many days a week?   "))
    # num_nights = int(input("How many nights agents will go to bar?   "))
    # num_agents = int(input("How many agents total?   "))
    # discount_rate = float(input("Please inpout discount-rate:   "))
    # reward_function = int(input("Which of the following reward-function\n"
    #                         "1:wonderful-life-0 2:wonderful-life-1  3:aristocratic"))
    days, num_nights, num_agents, discount_rate, reward_function = 7, 1, 63, 0.98, 3

    actions = []
    for i in range(days):
        action = [0 for j in range(days)]
        for k in range(num_nights):
            action[(i+k) % 7] = 1
        actions.append(action)

    # 创建所有agent对象
    agents = []
    for j in range(num_agents):
        agents.append(Agent(days, num_nights, actions))

    for m in range(100):
        for agent in agents:
            agent.take_random_action(actions)
        attendance = get_attendance(agents)
        for agent in agents:
            agent.learn(days, num_nights, discount_rate, reward_function, actions, attendance)
    return agents, reward_function, actions, attendance


# 计算每天出席的总人数
def get_attendance(agents):
    attendance = np.zeros(len(agents[0].last_action))
    for agent in agents:
        attendance += agent.last_action
    return attendance



def main(cicle):
    global days, num_nights, num_agents, discount_rate
    agents, reward_function, actions, attendance = setup()

    for i in range(cicle):
        for agent in agents:
            agent.take_action(actions)
        attendance = get_attendance(agents)
        for agent in agents:
            agent.learn(days, num_nights, discount_rate, reward_function, actions, attendance)

        print("第%s轮" %i)
        print(attendance)
        print(global_utility(attendance, num_nights))



main(8000)